[V1] Optimize block table transfer from CPU to GPU #11401
Conversation
Signed-off-by: Woosuk Kwon <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
| int* d_matrix_tgt = matrix_tgt.data_ptr<int>(); | ||
|
|
||
| // One thread block per row. | ||
| int blocks = n; |
There was a problem hiding this comment.
it seems this can easily oversubscribe GPU SMs.
| int length = matrix_diff[row_id * 2 + 1]; | ||
| int end = start + length; | ||
| int thread_idx = threadIdx.x; | ||
| for (int i = start + thread_idx; i < end; i += blockDim.x) { |
There was a problem hiding this comment.
most threads in the block would be idle, e.g. for decoding, there's only one or even no entry changes in the block table.
| self.block_table_diff_np[row_idx, 0] = start | ||
| # Move-and-append is not allowed. | ||
| assert self.block_table_diff_np[row_idx, 1] == 0 | ||
| self.block_table_diff_np[row_idx, 1] = num_blocks |
There was a problem hiding this comment.
for the non-uva case, we still need to keep track of the max-block-table-length, so that apply_diff only needs to copy max-block-table-length columns.
There was a problem hiding this comment.
Good point. The problem is, the memcpy API requires the data to be in contiguous memory space: https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#group__CUDART__MEMORY_1g85073372f776b4c4d5f89f7124b7bf79
So when the block table tensor has the shape [batch_size, max_model_len] and if we slice over the second dimension, then we have to call the memcpy API batch_size times instead of once.
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
WoosukKwon
requested review from
tlrmchlsmth,
robertgshaw2-neuralmagic,
njhill,
ywang96,
comaniac and
alexm-neuralmagic
as code owners
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
| int end = start + length; | ||
| int thread_idx = threadIdx.x; | ||
| for (int i = start + thread_idx; i < end; i += blockDim.x) { | ||
| int idx = row_offset + i; |
There was a problem hiding this comment.
Should row_offset and idx be int64_t? I.e. could they overflow an int32?
Signed-off-by: Woosuk Kwon <[email protected]>
Currently, the block table transfer from CPU to GPU could be expensive because we send the entire block table (
[batch_size, max_model_len // block_size]) every step. This PR optimizes the overhead by only sending the diffs from CPU to GPU, which is typically very small.The solution in this PR relies on CUDA unified virtual addressing, so may not work in some environments. In such a case, we fall back to the original implementation (copying the entire block table tensor).